# https://www.gairuo.com/p/pandas-wide-to-long
import pandas as pd
import numpy as np

np.random.seed(123)
df = pd.DataFrame(
    data={"A1970": {0: "a", 1: "b", 2: "c"},
          "A1980": {0: "d", 1: "e", 2: "f"},
          "B1970": {0: 2.5, 1: 1.2, 2: .7},
          "B1980": {0: 3.2, 1: 1.3, 2: .1},
          "X": dict(zip(range(3), np.random.randn(3)))
          }
)
df["id"] = df.index

print(df)
"""
  A1970 A1980  B1970  B1980         X  id
0     a     d    2.5    3.2 -1.085631   0
1     b     e    1.2    1.3  0.997345   1
2     c     f    0.7    0.1  0.282978   2
"""

ldf = pd.wide_to_long(
    df=df,  # DataFrame 数据，是一个 wide-format 数据
    stubnames=["A", "B"],  # str or list-like，存根（或者占坑、桩，stub）名称。宽格式变量假定以存根名称开头
    i="id",  # str or list-like，用作 id 变量的列
    j="year"  # str，子观察变量的名称，希望以长格式命名后缀
)
print(ldf)
"""
                X  A    B
id year                  
0  1970 -1.085631  a  2.5
1  1970  0.997345  b  1.2
2  1970  0.282978  c  0.7
0  1980 -1.085631  d  3.2
1  1980  0.997345  e  1.3
2  1980  0.282978  f  0.1
"""

# 具有多个id列 =================================
df = pd.DataFrame(
    data={
        'famid': [1, 1, 1, 2, 2, 2, 3, 3, 3],
        'birth': [1, 2, 3, 1, 2, 3, 1, 2, 3],
        'ht1': [2.8, 2.9, 2.2, 2, 1.8, 1.9, 2.2, 2.3, 2.1],
        'ht2': [3.4, 3.8, 2.9, 3.2, 2.8, 2.4, 3.3, 3.4, 2.9]
    }
)
print(df)
"""
   famid  birth  ht1  ht2
0      1      1  2.8  3.4
1      1      2  2.9  3.8
2      1      3  2.2  2.9
3      2      1  2.0  3.2
4      2      2  1.8  2.8
5      2      3  1.9  2.4
6      3      1  2.2  3.3
7      3      2  2.3  3.4
8      3      3  2.1  2.9
"""
l = pd.wide_to_long(
    df=df,
    stubnames="ht",
    i=["famid", "birth"],
    j="age"
)
print(l)
"""
                  ht
famid birth age     
1     1     1    2.8
            2    3.4
      2     1    2.9
            2    3.8
      3     1    2.2
            2    2.9
2     1     1    2.0
            2    3.2
      2     1    1.8
            2    2.8
      3     1    1.9
            2    2.4
3     1     1    2.2
            2    3.3
      2     1    2.3
            2    3.4
      3     1    2.1
            2    2.9
"""

# 从长到宽只需要创造性地使用“unstack”
w = l.unstack()
w.columns = w.columns.map("{0[0]}{0[1]}".format)  # 给列名加上前缀
w.reset_index()  # 重置索引，使数据帧的索引从0开始，同时将"famid"、"birth"列转化为普通列
print(w)

# 也会处理使用较少的列名
np.random.seed(0)
df = pd.DataFrame(
    data={
        'A(weekly)-2010': np.random.rand(3),
        'A(weekly)-2011': np.random.rand(3),
        'B(weekly)-2010': np.random.rand(3),
        'B(weekly)-2011': np.random.rand(3),
        'X': np.random.randint(3, size=3)
    }
)
df['id'] = df.index
print(df)

l = pd.wide_to_long(
    df=df,
    stubnames=["A(weekly)", "B(weekly)"],
    i="id",
    j="year",
    sep="-"
)
print(l)

# 如果我们有很多列，我们也可以使用正则表达式来查找stubnames并将该列表传递给 wide_to_long：
stubnames = sorted(
    set([match[0] for match in df.columns.str.findall(
        r'[A-B]\(.*\)').values if match != []])
)
print(list(stubnames))

# 以上所有示例都以整数作为后缀，可以将非整数作为后缀
df = pd.DataFrame(
    data={
        'famid': [1, 1, 1, 2, 2, 2, 3, 3, 3],
        'birth': [1, 2, 3, 1, 2, 3, 1, 2, 3],
        'ht_one': [2.8, 2.9, 2.2, 2, 1.8, 1.9, 2.2, 2.3, 2.1],
        'ht_two': [3.4, 3.8, 2.9, 3.2, 2.8, 2.4, 3.3, 3.4, 2.9]
    }
)
l = pd.wide_to_long(
    df=df,
    stubnames="ht",
    i=["famid","birth"],
    j="age",
    sep="_",
    suffix=r"\w+"
)
"""
参数suffix:
    str, default '\d+'，捕获所需后缀的正则表达式。 
    '\d+' 捕获数字后缀，可以指定没有数字的后缀否定字符类 '\D+'。 
    还可以进一步消除歧义后缀，例如，如果您的宽变量是 A-one 形式，B-2，..，并且您有一个不相关的列 A-评级，您可以忽略最后一个通过指定suffix='(!?one|two)'。 
    当所有后缀都是数字，它们被转换为 int64/float64
"""
print(l)
